# coding=utf-8
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import math
import collections

from pathlib import Path
from itertools import combinations, permutations
from imblearn.pipeline import make_pipeline
from imblearn.under_sampling import RandomUnderSampler
from numpy import interp
from sklearn import metrics
from sklearn.metrics import accuracy_score, confusion_matrix, multilabel_confusion_matrix, roc_curve, auc
from sklearn.model_selection import KFold
from sklearn.neighbors import KNeighborsClassifier


def cooccurrence_matrix(label, categories, fold):
    lbl_num = label.shape[1]

    # Compute cooccurrence matrix
    cooccurrence_matrix = np.dot(label.transpose(), label)
    # print('\ncooccurrence_matrix:\n{0}'.format(cooccurrence_matrix))

    # Compute the Jaccard coefficient
    jaccard_matrix = np.zeros(shape=(lbl_num, lbl_num))
    for i in range(lbl_num):
        for j in range(lbl_num):
            lbl_occur_union = cooccurrence_matrix[i][i] + cooccurrence_matrix[j][j] - cooccurrence_matrix[i][j]
            jaccard_matrix[i][j] = np.true_divide(cooccurrence_matrix[i][j], lbl_occur_union)
    # print('\njaccard_matrix:\n{0}'.format(jaccard_matrix))

    # Compute normalized similarities matrix
    norm_similarity_matrix = np.zeros(shape=(lbl_num, lbl_num))
    for i in range(lbl_num):
        quadratic_sum = 0
        for j in range(lbl_num):
            quadratic_sum = quadratic_sum + math.pow(jaccard_matrix[i][j], 2)
        norm_numerator = math.sqrt(quadratic_sum)
        norm_similarity_matrix[i, :] = np.true_divide(jaccard_matrix[i, :], norm_numerator)
    # print('\nnorm_similarity_matrix:\n{0}'.format(norm_similarity_matrix))

    # draw
    df_cm_per = pd.DataFrame(jaccard_matrix, index=categories, columns=categories)
    fig = plt.figure(figsize=(12, 8))

    # annot: If True, write the data value in each cell. annot_kws={"fontsize": "small/medium/large"}
    Heatmap = sns.heatmap(df_cm_per, annot=True, fmt=".2f", annot_kws={"fontsize": "medium"}, cmap='OrRd',
                          linewidths=0.5,
                          linecolor='white', square=True, mask=False, cbar_kws={"orientation": "vertical"}, cbar=True)

    plt.xticks(rotation=90, fontsize=15)
    plt.yticks(fontsize=15)
    plt.tight_layout()  # fix the xlabels cutoff. X-axis text is too lengthy
    plt.savefig("Jaccard Matrix for fold " + str(fold+1) + ".pdf")
    # plt.show()
    plt.close()

    return norm_similarity_matrix


def get_evaluation_single(true_value, prediction_value):
    cnf_matrix = confusion_matrix(true_value, prediction_value)
    # print(cnf_matrix)

    FP = cnf_matrix.sum(axis=0) - np.diag(cnf_matrix)
    FN = cnf_matrix.sum(axis=1) - np.diag(cnf_matrix)
    TP = np.diag(cnf_matrix)
    TN = cnf_matrix.sum() - (FP + FN + TP)
    FP = FP.astype(float)
    FN = FN.astype(float)
    TP = TP.astype(float)
    TN = TN.astype(float)

    ACC = accuracy_score(true_value, prediction_value)

    # Sensitivity, hit rate, recall, or true positive rate
    TPR = metrics.recall_score(true_value, prediction_value, average='macro')

    # Specificity or true negative rate
    TNR = np.mean(TN / (TN + FP))

    # Precision or positive predictive value
    F1_score = metrics.f1_score(true_value, prediction_value, average='macro')

    return ACC, TPR, TNR, F1_score, cnf_matrix


def get_evaluation_multi(true_value, prediction_value):
    cnf_matrix = multilabel_confusion_matrix(true_value, prediction_value)
    # print(cnf_matrix)

    TN = cnf_matrix[:, 0, 0]
    TP = cnf_matrix[:, 1, 1]
    FN = cnf_matrix[:, 1, 0]
    FP = cnf_matrix[:, 0, 1]
    FP = FP.astype(float)
    FN = FN.astype(float)
    TP = TP.astype(float)
    TN = TN.astype(float)

    ACC = accuracy_score(true_value, prediction_value)

    # Sensitivity, hit rate, recall, or true positive rate
    TPR = metrics.recall_score(true_value, prediction_value, average='macro')

    # Specificity or true negative rate
    TNR = np.mean(TN / (TN + FP))

    # Precision or positive predictive value
    F1_score = metrics.f1_score(true_value, prediction_value, average='macro')

    return ACC, TPR, TNR, F1_score


def pairwise_accuracy(true_value, prediction_value):
    f1 = 1.0 * ((true_value > 0) & (prediction_value > 0)).sum(axis=1)
    f2 = 1.0 * ((true_value > 0) | (prediction_value > 0)).sum(axis=1)
    f1[f2 == 0] = 1.0
    f1[f2 > 0] /= f2[f2 > 0]
    return f1


def jackknife_resampling(data):
    """
    Performs jackknife resampling on numpy arrays.
    Jackknife resampling is a technique to generate 'n' deterministic samples
    of size 'n-1' from a measured sample of size 'n'. Basically, the i-th
    sample, (1<=i<=n), is generated by means of removing the i-th measurement
    of the original sample. Like the bootstrap resampling, this statistical
    technique finds applications in estimating variance, bias, and confidence
    intervals.
    Parameters
    ----------
    data : numpy.ndarray
        Original sample (1-D array) from which the jackknife resamples will be
        generated.
    Returns
    -------
    resamples : numpy.ndarray
        The i-th row is the i-th jackknife sample, i.e., the original sample
        with the i-th measurement deleted.
    References
    ----------
    .. [1] McIntosh, Avery. "The Jackknife Estimation Method".
        <http://people.bu.edu/aimcinto/jackknife.pdf>
    .. [2] Efron, Bradley. "The Jackknife, the Bootstrap, and other
        Resampling Plans". Technical Report No. 63, Division of Biostatistics,
        Stanford University, December, 1980.
    .. [3] Cowles, Kate. "Computing in Statistics: The Jackknife, Lecture 11".
        <http://homepage.stat.uiowa.edu/~kcowles/s166_2009/lect11.pdf>.
        September, 2009.
    .. [4] Jackknife resampling <https://www.stat.berkeley.edu/~hhuang/STAT152/Jackknife-Bootstrap.pdf>
    .. [5] Jackknife resampling <https://en.wikipedia.org/wiki/Jackknife_resampling>
    """
    n = data.shape[0]
    if n <= 0:
        raise ValueError("data must contain at least one measurement.")

    resamples = np.empty([n, n-1])

    for i in range(n):
        resamples[i] = np.delete(data, i)

    return resamples


def delete_group_jackknife_resampling(data, n_delete=1):
    """
    An alternate the jackknife method of deleting one observation at a time is to delete d.
    this function only delete d in order, does not fulfill a combination C(n,d).
    this function can: verify whether jackknife_resampling function is running

    TODO:
    C(n,d）
    delete_d_jackknife_resampling
    It is natural to extend the delete-1 jackknife by omitting more than one observation.
    Instead of leaving out one observation at a time, we leave out d observations.
    Therefore, the size of a delete-d jackknife sample is (n-d), and there are C(n,d) jackknife samples.
    In the simplest case jackknife resampling is generated by sequentially deleting single cases
    from the original sample (delete‐one jackknife).
    A more generalized jackknife technique uses resampling that is based on multiple case deletion (delete‐d jackknife).

    References
    ----------
        https://www.stat.berkeley.edu/~hhuang/STAT152/Jackknife-Bootstrap.pdf
    """
    n = data.shape[0]  # number of data points
    if n <= 0:
        raise ValueError("data must contain at least one measurement.")

    # combins = [c for c in combinations(range(n), n_delete)]  # C(n,d)
    n_sub = n // n_delete  # number of data subset(model)(In Python 3, a//b floor division, a/b "true" (float) division)
    start = 0  # start of current block of data
    resamples = np.empty([n_sub, n-n_delete])

    for i in range(0, n_sub):
        end = np.minimum(start + n_delete, n)
        r = range(start, end)

        resamples[i] = np.delete(data, r)
        start = end

    return resamples


def prediction_model(X, y, categories, threshold, organ_choice):
    # K-Folds cross-validator
    fold_num = 5
    kf = KFold(n_splits=fold_num)
    label_num = y.shape[1]

    ACC = []
    SPE = []
    TPR = []
    F1_score = []

    Label_Acc = [[] for i in range(label_num)]  # expect to access the performance of EACH LABEL
    Average_Label_Accuracy = []
    mean_Label_Acc = []
    Pairwise_Acc = []

    tprs = []
    aucs = []
    mean_fpr = np.linspace(0, 1, 100)

    integrative_prediction_score_save = [[] for i in range(fold_num)]

    # define undersample strategy
    rus = RandomUnderSampler(random_state=222, sampling_strategy='majority')  # resample the majority class

    fold = 0
    for train_index, test_index in kf.split(X, y):
        print("K-fold cross validation: fold " + str(fold+1))
        # print("TRAIN:", train_index, "Validation:", validation_index)
        # Try to change the data structure from pandas to numpy instead, using .values
        X_train, X_test = X.values[train_index], X.values[test_index]
        Y_train, Y_test = y.values[train_index], y.values[test_index]

        # occurrence matrix that describes the presentation states of training samples
        normalized_similarity_matrix = cooccurrence_matrix(Y_train, categories, fold)

        clf = KNeighborsClassifier(n_neighbors=5)

        tst_sample_num = X_test.shape[0]

        prediction_score = np.zeros([label_num, tst_sample_num])
        for lbl in range(label_num):
            # # fit and apply the transform
            # X_train_rus, Y_train_rus = rus.fit_resample(X_train, Y_train[:, lbl])
            # print("Y_train_rus" + str(Y_train_rus))
            # print('Resampled dataset shape %s' % collections.Counter(Y_train_rus))

            # to get sample indices from RandomUnderSampler in imblearn
            # As the undersampling is solely based on the y_vector,
            # one can add a counter-variable instead of the the x-vector/array
            counter = range(0, len(Y_train[:, lbl]))
            counter = np.array(counter).reshape(-1, 1)
            train_index_rus, Y_train_rus = rus.fit_resample(counter, Y_train[:, lbl])
            print('Resampled dataset shape %s' % collections.Counter(Y_train_rus))

            train_index_res = jackknife_resampling(train_index_rus)
            # train_index_res = delete_group_jackknife_resampling(train_index_rus, 5)
            # print("train_index_res:"+str(train_index_res))  # train set index after jackknife resample
            n_sub = train_index_res.shape[0]

            sub_prediction_score = np.zeros([tst_sample_num, n_sub])  # n sub prediction models for each label
            for idx in range(n_sub):
                idx_res = train_index_res[idx].astype(np.int)
                X_train_res = X.values[idx_res]
                Y_train_res = y.values[idx_res]

                clf.fit(X_train_res, Y_train_res[:, lbl])  # the i-th label
                y_score = clf.predict_proba(X_test)
                # print("y_score:" + str(y_score))
                y_pred = clf.predict(X_test)
                # print("y_pred:" + str(y_pred))

                # store sub prediction score
                for sample in range(tst_sample_num):
                    # sub_prediction_score[sample][idx] = y_score[sample][1]  # higher than a threshold (would occur)
                    sub_prediction_score[sample][idx] = y_pred[sample]
            # print("sub_prediction_score:" + str(sub_prediction_score))

            for sample2 in range(tst_sample_num):
                prediction_score[lbl][sample2] = np.mean(sub_prediction_score[sample2, :])
            # print("prediction_score(lbl):" + str(prediction_score))
        print("prediction_score(model):" + str(prediction_score))

        # without normalized similarity matrix (without weight)
        # for row in range(0, prediction_score.shape[0]):
        #     for col in range(0, prediction_score.shape[1]):
        #         if prediction_score[row, col] > threshold:
        #             prediction_score[row, col] = 1
        #         else:
        #             prediction_score[row, col] = 0
        # print(prediction_score)

        # with normalized similarity matrix (with weight)
        integrative_prediction_score = np.dot(normalized_similarity_matrix, prediction_score)  # matrix multiplication
        print("integrative_prediction_score:" + str(integrative_prediction_score))
        integrative_predict = np.zeros((integrative_prediction_score.shape[0], integrative_prediction_score.shape[1]))
        for row in range(0, integrative_prediction_score.shape[0]):
            for col in range(0, integrative_prediction_score.shape[1]):
                if integrative_prediction_score[row, col] > threshold:
                    integrative_predict[row, col] = 1
                else:
                    integrative_predict[row, col] = 0
        # print("integrative_predict:" + str(integrative_predict))
        # print("Y_test:" + str(Y_test))

        # save each fold integrative_prediction_score
        integrative_prediction_score_save[fold] = integrative_prediction_score  # add each fold in a list
        # print("integrative_prediction_score_save:"+str(integrative_prediction_score_save))

        # Compute ROC curve and area the curve
        fpr, tpr, thresholds = roc_curve(Y_test.ravel(), integrative_predict.ravel())
        tprs.append(interp(mean_fpr, fpr, tpr))
        tprs[-1][0] = 0.0
        roc_auc = auc(fpr, tpr)
        aucs.append(roc_auc)

        ACC_tst, TPR_tst, SPE_tst, F1_score_tst = get_evaluation_multi(Y_test, integrative_predict.T)
        Pairwise_Acc_tst = np.mean(pairwise_accuracy(Y_test, integrative_predict.T))
        ACC.append(ACC_tst)
        TPR.append(TPR_tst)
        SPE.append(SPE_tst)
        F1_score.append(F1_score_tst)
        Pairwise_Acc.append(Pairwise_Acc_tst)

        each_label_value = []  # each_label_value[0] represents the first label true value
        each_estimators_predict = integrative_predict.tolist()

        # transposition, save each label true value in a list
        for i in range(len(Y_test[0])):  # row num
            t = []
            for j in range(len(Y_test)):
                t.append(Y_test[j][i])
            each_label_value.append(t)

        lbl_acc_sum = 0
        for i in range(0, label_num, 1):
            lbl_pred = each_estimators_predict[i]  # the i-th label prediction, each label predict value
            lbl_acc = accuracy_score(each_label_value[i], lbl_pred)  # each label accuracy in one fold
            Label_Acc[i].append(lbl_acc)
            lbl_acc_sum = lbl_acc_sum + lbl_acc

        Average_Label_Accuracy.append(lbl_acc_sum / label_num)

        fold = fold + 1

    # Draw Diagonal Lines
    plt.plot((0, 1), (0, 1), c='#808080', lw=1, ls='--', alpha=0.7)
    # calculate mean auc
    mean_tpr = np.mean(tprs, axis=0)
    mean_tpr[-1] = 1.0
    mean_auc = auc(mean_fpr, mean_tpr)
    plt.plot(mean_fpr, mean_tpr, color='Blue',
             label=r'Mean ROC(Area=%0.4f)' % mean_auc, lw=1, alpha=.8)

    # k-fold mean value of label accuracy
    for i in range(0, label_num, 1):
        mean_Label_Acc.append(np.mean(Label_Acc[i]))

    print("--------------test set result------------------")
    print("test:" + str(ACC))
    print(organ_choice, 'ACC:', np.mean(ACC), 'TPR:', np.mean(TPR), 'SPE', np.mean(SPE), 'F1_score:', np.mean(F1_score),
          'AUC:', mean_auc, 'Pairwise_Acc:', np.mean(Pairwise_Acc), 'Ave_Label_Acc:', np.mean(Average_Label_Accuracy))

    # Writing CSV Files
    columns_name = ['organ_choice', 'ACC', 'TPR', 'spe', 'f1', 'auc', 'pairAcc', 'aveLabAcc',
                    'label1', 'label2', 'label3', 'label4', 'label5', 'label6',
                    'label7', 'label8']
    list_1 = [organ_choice, np.mean(ACC), np.mean(TPR), np.mean(SPE), np.mean(F1_score), mean_auc,
              np.mean(Pairwise_Acc), np.mean(Average_Label_Accuracy)] + mean_Label_Acc
    pd_data = pd.DataFrame(np.array(list_1).reshape(1, 16), columns=columns_name)
    output_path = "./" + organ_choice + "_integrative_model_test_result.csv"
    pd_data.to_csv(output_path, index=False)

    print("integrative_prediction_score_save:" + str(integrative_prediction_score_save))
    return integrative_prediction_score_save


def draw_roc_individual_model(X, y, categories, integrative_prediction_score_fold_save, organ_choice):
    # K-Folds cross-validator
    kf = KFold(n_splits=5)
    label_num = y.shape[1]

    color_list = ['ForestGreen', 'Navy', 'SaddleBrown', 'DarkCyan', 'DarkOrange', 'DarkMagenta', 'Olive', 'DarkGray',
                  'Tomato', 'Tan', 'plum', 'CadetBlue']

    for lbl in range(label_num):
        yy = y.iloc[:, lbl]  # the i-th label
        tprs = []
        mean_fpr = np.linspace(0, 1, 100)

        fold = 0
        for train_index, test_index in kf.split(X, yy):
            print("fold:"+str(fold+1))
            X_train, X_test = X.values[train_index], X.values[test_index]
            Y_train, Y_test = yy.values[train_index], yy.values[test_index]

            y_score = integrative_prediction_score_fold_save[fold][lbl]
            print(y_score)
            # print(Y_test.shape)
            # print(y_score.shape)

            # Compute ROC curve and area the curve
            fpr, tpr, thresholds = roc_curve(Y_test, y_score)
            tprs.append(interp(mean_fpr, fpr, tpr))
            tprs[-1][0] = 0.0

            fold = fold + 1

        mean_tpr = np.mean(tprs, axis=0)
        mean_tpr[-1] = 1.0
        mean_auc = auc(mean_fpr, mean_tpr)
        plt.plot(mean_fpr, mean_tpr, color=color_list[lbl],
                 label=r'%s (Area=%0.4f)' % (categories[lbl], mean_auc), lw=1, alpha=.8)

    # Draw Diagonal Lines
    plt.plot((0, 1), (0, 1), c='#808080', lw=1, ls='--', alpha=0.7)
    # Plot them on canvas and give a name to x-axis and y-axis
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    # plt.legend(loc='lower right')
    plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))  # Put a legend to the right of the current axis
    plt.savefig("integrative model roc all " + organ_choice + ".pdf", bbox_inches='tight')
    plt.show()


def main():
    # open the dataset you selected (check if it exists)
    # organ = "liver"
    organ = "kidney"
    # data_path = Path('dataset/liver-MLFS-sorted-feature.csv')
    data_path = Path('dataset/kidney-MLFS-sorted-feature.csv')

    if not data_path.exists():
        print("Oops,the file: " + str(data_path) + " ,doesn't exist!")
    else:
        print("Yay,the file: " + str(data_path) + " ,exists!")

    # read it into dataframe
    dataset = pd.read_csv(data_path)
    # print(dataset.shape)

    # extract labels from data
    features = dataset.iloc[:, 8:]  # features
    labels = dataset.iloc[:, :8]  # labels

    # keep label names in a list
    categories = labels.columns.tolist()
    if organ == "liver":
        categories = ['Cellular infiltration', 'Eosinophilic change', 'Hypertrophy', 'Increased mitosis',
                      'NOS lesion', 'Microgranuloma', 'Necrosis', 'Hepatodiaphragmatic nodule',
                      'Kupffer cell proliferation', 'Single cell necrosis', 'Swelling', 'Cytoplasmic vacuolization']
    elif organ == "kidney":
        categories = ['Hyaline cast', "Lymphocyte cellular infiltration", 'Basophilic change', 'Cyst',
                      'Dilatation', 'Cystic dilatation', 'Necrosis', 'Regeneration']

    # Normalise the data
    features = (features - features.min()) / (features.max() - features.min())  # min-max normalization

    threshold = 0.5
    features_select = features.iloc[:, 0: 200]  # feature subsets
    integrative_prediction_score_save = prediction_model(features_select, labels, categories, threshold, organ)
    draw_roc_individual_model(features_select, labels, categories, integrative_prediction_score_save, organ)


if __name__ == "__main__":
    main()

